#-*- coding: utf-8 -*-
#
# Copyright 2013 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

import math
import numpy as np
import control.matlab

import pat.vehicles.rotorcraft.dynamic_model_synth as dm
import pat.autopilot as ap
import pat.utils as pu
import pat.trajectory as pt

class Autopilot(ap.Autopilot):
    def __init__(self, _args=None, _dm=None, _time=None):
        ap.Autopilot.__init__(self)
        Xe, Ue = _dm.trim(); A, B = _dm.get_jacobian(Xe, Ue)
        for law in [TestLawLQR, LawGo]:
            self.add_mode(law(_dm, Xe, Ue, A, B, time))


class TestLawLQR(ap.LQRFeedback):
    def __init__(self, dm, Xe, Ue, A, B, time):
        #                  x       y      z     u       v      w    phi   theta psi     p       q       r   om
        self.Q = np.diag([0.005, 0.005, 0.1, 0.005, 0.005, 0.005,  0.05, 0.05, 0.1, 0.001,  0.001,  0.001, 0.01])
        #                 lat   lon   rud   col   thr 
        self.R = np.diag([0.01, 0.01, 0.05, 0.05, 0.1])
        (K, X, E) = control.matlab.lqr(A, B, self.Q, self.R)
        refs = [Step_x(Xe, Ue), Step_z(Xe, Ue), Step_psi(Xe, Ue)]
        ap.LQRFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, -K, "PositionLQR", refs)


class LawGo(ap.LQRFeedback):
    def __init__(self, dm, Xe, Ue, A, B, time):
        #                  x       y      z     u       v      w    phi   theta psi     p       q       r   om
        self.Q = np.diag([0.005, 0.005, 0.1, 0.005, 0.005, 0.005,  0.05, 0.05, 0.1, 0.001,  0.001,  0.001, 0.01])
        #                 lat   lon   rud   col   thr 
        self.R = np.diag([0.01, 0.01, 0.05, 0.05, 0.1])
        (K, X, E) = control.matlab.lqr(A, B, self.Q, self.R)
        refs = [RefGo(Xe, Ue, "RefGo", time), RefGo2(Xe, Ue, "RefGo2")]
        ap.LQRFeedback.__init__(self, dm.sv_size, dm.iv_size, Xe, Ue, -K, "LawGo", refs)



class Step_x(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, dm.sv_x, "Step x", ampl=2.)

class Step_z(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, dm.sv_z, "Step z", ampl=1.)

class Step_psi(ap.StepReference):
    def __init__(self, Xr, Ur): ap.StepReference.__init__(self, Xr, Ur, dm.sv_psi, "Step psi", ampl=pu.rad_of_deg(45.))



class RefGo(ap.Reference):

    def __init__(self, Xr, Ur, name, tr=None):
        ap.Reference.__init__(self, Xr, Ur, name)
        self.Xr1 = Xr
        self.tr = pt.Tlti(dt=0.01)
        self.set_param(p=Xr, t0=1.)
        
    def set_param(self, p, t0=0):
        print "hello"
        self.X1 = p
        self.v = 0.5
        self.a = 1.
        self.t0 = t0
        self.dx = p[dm.sv_x]-self.trim[0][dm.sv_x]
        self.dt = self.dx/self.v
        self.dt_a = 2.
        print self.dt
        self.period = 15
    
    def get(self, time):
        Xr1 = np.array(self.trim[0])
        if time > self.t0:
            a = math.fmod(time-self.t0, self.period)
            if a >= self.period/2. :
                Xr = self.tr.get(time, U=5)
            else:
                Xr = self.tr.get(time, U=-5)
            self.Xr1[dm.sv_x] = Xr[0]
            self.Xr1[dm.sv_u] = Xr[1]
            self.Xr1[dm.sv_theta] = -math.atan2(Xr[2], 9.81)
        Ur1 = np.array(self.trim[1])
        
        return self.Xr1, Ur1 


class RefGo2(ap.Reference):
    
    def __init__(self, Xr, Ur, name="Go2", period=20, dt=5.):
        ap.Reference.__init__(self, Xr, Ur, name)
        self.period = period
        self.dt = dt
        self.t0 = 0.
        self.v = 2.
        
    def get(self, time):
        self.Xr1 = np.array(self.trim[0])
        if time >= self.t0: 
            a = math.fmod(time-self.t0, self.period)
            if a < self.dt :
                self.Xr1[dm.sv_x] += self.v*a -5
                self.Xr1[dm.sv_u]  = self.v
            elif a<self.period/2.:
                self.Xr1[dm.sv_x] += self.v*self.dt -5
            elif a<self.period/2.+self.dt:
                self.Xr1[dm.sv_x] += -self.v*(a-self.period/2-self.dt) -5
                self.Xr1[dm.sv_u]  = -self.v
            else:
               self.Xr1[dm.sv_x] += -5 
            
        Ur1 = np.array(self.trim[1])
        
        return self.Xr1, Ur1 
        
    def set_param(self, p, t0=0):
        pass
